import asyncio
from typing import Callable, List, Any, Coroutine, Dict, Optional, Union
from enum import Enum
from dataclasses import dataclass
from queue import PriorityQueue
import time
import logging
from asyncio import TimeoutError

class TaskStatus(Enum):
    PENDING = "待执行"
    RUNNING = "执行中"
    COMPLETED = "已完成"
    FAILED = "失败"
    CANCELLED = "已取消"
    TIMEOUT = "超时"

class TaskPriority(Enum):
    HIGH = 1
    MEDIUM = 2
    LOW = 3

@dataclass
class TaskInfo:
    id: str
    coro: Coroutine
    priority: TaskPriority
    status: TaskStatus
    timeout: Optional[float]
    start_time: float = 0
    end_time: float = 0
    result: Any = None
    error: Optional[Exception] = None

class AsyncManager:
    def __init__(self, max_concurrent: int = 3, default_timeout: Optional[float] = None):
        self.max_concurrent = max_concurrent
        self.semaphore = asyncio.Semaphore(max_concurrent)
        self.task_queue = PriorityQueue()
        self.tasks: Dict[str, TaskInfo] = {}
        self.loop = asyncio.get_event_loop()
        self._task_id_counter = 0
        self.default_timeout = default_timeout

    def _generate_task_id(self) -> str:
        self._task_id_counter += 1
        return f"task_{self._task_id_counter}"

    async def execute_with_limit(self, task_info: TaskInfo) -> Any:
        """使用信号量限制并发并监控执行状态"""
        async with self.semaphore:
            try:
                task_info.status = TaskStatus.RUNNING
                task_info.start_time = time.time()
                
                timeout = task_info.timeout or self.default_timeout
                if timeout:
                    try:
                        result = await asyncio.wait_for(task_info.coro, timeout=timeout)
                    except TimeoutError:
                        task_info.status = TaskStatus.TIMEOUT
                        task_info.error = TimeoutError(f"任务执行超过 {timeout} 秒")
                        raise task_info.error
                else:
                    result = await task_info.coro
                
                task_info.status = TaskStatus.COMPLETED
                task_info.result = result
                return result
                
            except TimeoutError:
                task_info.status = TaskStatus.TIMEOUT
                raise
            except Exception as e:
                task_info.status = TaskStatus.FAILED
                task_info.error = e
                raise
            finally:
                task_info.end_time = time.time()
                self._update_task_stats(task_info)

    def _update_task_stats(self, task_info: TaskInfo):
        """更新任务统计信息"""
        duration = task_info.end_time - task_info.start_time
        logging.info(f"任务 {task_info.id} - 状态: {task_info.status.value}")
        logging.info(f"执行时间: {duration:.2f}秒")
        if task_info.error:
            logging.error(f"错误信息: {str(task_info.error)}")

    async def add_task(self, 
                      coro: Coroutine, 
                      priority: TaskPriority = TaskPriority.MEDIUM,
                      timeout: Optional[float] = None) -> str:
        """添加带优先级和超时的异步任务"""
        task_id = self._generate_task_id()
        task_info = TaskInfo(
            id=task_id,
            coro=coro,
            priority=priority,
            status=TaskStatus.PENDING,
            timeout=timeout
        )
        self.tasks[task_id] = task_info
        self.task_queue.put((priority.value, task_id))
        return task_id

    async def process_queue(self, total_timeout: Optional[float] = None):
        """处理任务队列，可设置总体超时时间"""
        try:
            if total_timeout:
                async with asyncio.timeout(total_timeout):
                    await self._process_queue_internal()
            else:
                await self._process_queue_internal()
        except TimeoutError:
            # 处理总体超时
            remaining_tasks = list(self.task_queue.queue)
            for _, task_id in remaining_tasks:
                task_info = self.tasks[task_id]
                if task_info.status == TaskStatus.PENDING:
                    task_info.status = TaskStatus.TIMEOUT
            logging.error(f"队列处理总体超时 ({total_timeout}秒)")
            raise

    async def _process_queue_internal(self):
        """内部队列处理逻辑"""
        while not self.task_queue.empty():
            _, task_id = self.task_queue.get()
            task_info = self.tasks[task_id]
            try:
                await self.execute_with_limit(task_info)
            except TimeoutError:
                logging.error(f"任务 {task_id} 超时")
            except Exception as e:
                logging.error(f"任务 {task_id} 执行失败: {str(e)}")

    def get_task_status(self, task_id: str) -> Optional[TaskInfo]:
        """获取任务状态"""
        return self.tasks.get(task_id)

    def get_all_tasks_status(self) -> Dict[str, TaskInfo]:
        """获取所有任务状态"""
        return self.tasks.copy()

    def get_stats(self) -> Dict[str, int]:
        """获取任务统计信息"""
        stats = {status.value: 0 for status in TaskStatus}
        for task in self.tasks.values():
            stats[task.status.value] += 1
        return stats

    async def cancel_task(self, task_id: str) -> bool:
        """取消任务"""
        if task_id in self.tasks:
            task_info = self.tasks[task_id]
            if task_info.status == TaskStatus.PENDING:
                task_info.status = TaskStatus.CANCELLED
                return True
        return False 